using DataFrames
using QuadGK
using Distributions
using JLD
using Optim
using ForwardDiff
using CSV
using LinearAlgebra
using SpecialFunctions
using Random
using DelimitedFiles
using SparseArrays


clearconsole()

#-------------------------------------------------------------
# include Functions
#-------------------------------------------------------------
#cd("$(pwd())/Dropbox/Heterogeneity/Software/KS_Simulation/")
readDir = "$(pwd())/Functions/"
include(readDir *"vech.jl");
#include(readDir *"VAR_Procedures.jl");
include(readDir *"VAR_MoreLags_Procedures.jl");
include(readDir *"Loaddata.jl");

#-------------------------------------------------------------
# choose specification files
#-------------------------------------------------------------
nDataSel    = "2"
nVARSpec    = "4P" # with 1= percentiles, and  2= gini, frac
nVARMCMCSpec= "1"
specDir   = "$(pwd())/SpecFiles/"
include(specDir * "/AltVARspec" * nVARSpec * ".jl")
include(specDir * "/AltVARMCMCspec" * nVARMCMCSpec * ".jl")

#-------------------------------------------------------------
# load aggregate and alt data
#-------------------------------------------------------------
dataDir  = "$(pwd())/data/Para" * nDataSel *"/"
agg_data = loadaggdata(1,Tend)
n_agg    = size(agg_data)[2]
percentiles_data, vec_percs, gini_coef_data, probMass_below1_data, probMass_above4_data = loadaltdata(1,Tend)

#-------------------------------------------------------------
# Combine agg and alt data
#-------------------------------------------------------------
if nSpecStr == "Pctl"
    data    = [ agg_data percentiles_data ]
    n_cross = size(percentiles_data)[2]
else
    data    = [ agg_data gini_coef_data probMass_above4_data]
    n_cross = 2
end

# adjust sample
data_adj = data[Tdrop-nlags+1:end,:]
n = n_agg + n_cross
p = nlags

# OLS estimation
YY, XX, PHIols, Sols, SIGMAols = data_to_YY(data_adj, nlags, nlags)

#-------------------------------------------------------------
# Generate prior
#-------------------------------------------------------------
igammadf = size(YY)[2]+sigdf
prior    = prior_ACP_stru(YY, SIGMAols, n_agg, igammadf, nlags, lambda)

#-------------------------------------------------------------
# Run Posterior Sampler
#-------------------------------------------------------------

println("")
println("Bayesian estimation of VAR: Gibbs Sampling... ")
println("")

store_alp, store_beta, store_Sig = sample_ThetaSig(YY,XX,PHIols,Sols,SIGMAols,nlags,prior,nsim)
store_Btilde, store_Sigpdraw     = getReducedForm(store_alp,store_beta,store_Sig)

# Initialize matrices for collecting draws from posterior
store_Bpdraw             = zeros(nsim,n*p,n)
store_Sigtrpdraw         = zeros(nsim,n,n)

for jj = 1:nsim

    Sigpdraw_j = Symmetric(store_Sigpdraw[jj,:,:])
    L = cholesky(Sigpdraw_j).L
    store_Sigtrpdraw[jj,:,:] = L

    Btilde = reshape(store_Btilde[jj,:], n*p,n)
    global store_Bpdraw[jj,:,:] = Btilde

end

store_Bpdraw             = store_Bpdraw[nburn+1:nsim,:,:];
store_Sigpdraw           = store_Sigpdraw[nburn+1:nsim,:,:];
store_Sigtrpdraw         = store_Sigtrpdraw[nburn+1:nsim,:,:];


#-------------------------------------------------------------
# Save draws
#-------------------------------------------------------------

sName    = "AltVAR" * nVARSpec * "_MCMC" * nVARMCMCSpec * "_" * nSpecStr
savedir  = "$(pwd())/results/Para" *nDataSel * "/" * sName *"/";
try mkdir(savedir) catch; end
save(savedir * sName* "_PostDraws.jld", "Bpdraw", store_Bpdraw, "Sigpdraw", store_Sigpdraw, "Sigtrpdraw", store_Sigtrpdraw)

#-------------------------------------------------------------
# MISC ITEMS
#-------------------------------------------------------------

Bpmean       = dropdims(mean(store_Bpdraw,dims=1),dims=1)
Sigpmean     = dropdims(mean(store_Sigpdraw,dims=1),dims=1)
Sigtrpmean   = dropdims(mean(store_Sigtrpdraw,dims=1),dims=1)

save(savedir * sName* "_PostMeans.jld", "Bpmean", Bpmean, "Sigpmean", Sigpmean, "Sigtrpmean", Sigtrpmean)

# save OLS estiamtes and posterior means as CSV files
CSV.write(savedir * sName * "_PHIols.csv", DataFrame(PHIols,:auto))
CSV.write(savedir * sName * "_SIGMAols.csv", DataFrame(SIGMAols,:auto))
CSV.write(savedir * sName * "_Bpmean.csv", DataFrame(Bpmean,:auto))
CSV.write(savedir * sName * "_Sigpmean.csv", DataFrame(Sigpmean,:auto))
CSV.write(savedir * sName * "_Sigtrpmean.csv", DataFrame(Sigtrpmean,:auto))
